// 导入鸢尾花数据和类别信息
import { getIrisData, IRIS_CLASSES } from './data.js'
// 导入TensorFlow.js库
import * as tf from '@tensorflow/tfjs'
// 导入TensorFlow.js可视化库
import * as tfvis from '@tensorflow/tfjs-vis'

// 页面加载完成后执行
window.onload = async () => {
    // 更新状态显示，提示正在加载和训练模型
    updateResultDisplay('模型正在加载和训练中，请稍候...');
    
    try {
        // 获取鸢尾花数据，划分训练集和测试集（测试集比例0.15）
        const [xTrain, yTrain, xTest, yTest] = getIrisData(0.15)
        console.log('数据加载完成', xTrain.shape, yTrain.shape, xTest.shape, yTest.shape)
        console.log('鸢尾花类别:', IRIS_CLASSES)
        
        // 创建序列模型
        const model = tf.sequential()
        
        // 添加第一层隐藏层：10个神经元，输入形状为4（对应4个特征），激活函数为sigmoid
        model.add(tf.layers.dense({
            units: 10,
            inputShape: [4],
            activation: 'sigmoid'
        }))
        
        // 添加输出层：3个神经元（对应3种类别），激活函数为softmax
        model.add(tf.layers.dense({
            units: 3,
            activation: 'softmax'
        }))
        
        // 编译模型：损失函数为交叉熵，优化器为Adam，学习率0.1，评估指标为准确率
        model.compile({
            loss: 'categoricalCrossentropy',
            optimizer: tf.train.adam(0.1),
            metrics: ['accuracy']
        })
        
        // 训练模型：100个 epochs，使用测试集作为验证数据，同时显示训练可视化
        console.log('开始训练模型...');
        await model.fit(xTrain, yTrain, {
            epochs: 100,
            validationData: [xTest, yTest],
            callbacks: tfvis.show.fitCallbacks(
                { name: '训练效果'}, 
                ['loss', 'val_loss', 'acc', 'val_acc'],
                {callbacks: ['onEpochEnd']}
            )
        })
        
        // 训练完成，更新状态显示
        updateResultDisplay('模型训练完成，可以开始预测了！');
        console.log('模型训练完成');
        
        // 定义预测函数，供页面调用
        window.predict = (form) => {
            // 从表单获取输入值并转换为浮点数
            const a = parseFloat(form.a.value)
            const b = parseFloat(form.b.value)
            const c = parseFloat(form.c.value)
            const d = parseFloat(form.d.value)
            
            // 简单的输入验证
            if (isNaN(a) || isNaN(b) || isNaN(c) || isNaN(d)) {
                updateResultDisplay('请输入有效的数值！', 'error');
                return;
            }
            
            // 创建张量作为模型输入
            const input = tf.tensor([[a, b, c, d]])
            
            // 使用tf.tidy清理中间张量，避免内存泄漏
            tf.tidy(() => {
                // 进行预测
                const pred = model.predict(input)
                // 获取预测概率和类别索引
                const predictionData = pred.dataSync();
                const classIndex = pred.argMax(1).dataSync()[0];
                const className = IRIS_CLASSES[classIndex];
                
                // 格式化预测结果显示
                let resultText = `<strong>预测结果：${className}</strong><br>`;
                resultText += '<p><strong>各类别概率：</strong></p><ul>';
                
                // 显示每个类别的概率
                IRIS_CLASSES.forEach((class_, index) => {
                    const probability = (predictionData[index] * 100).toFixed(2);
                    resultText += `<li>${class_}: ${probability}%</li>`;
                });
                
                resultText += '</ul>';
                
                // 在页面上显示预测结果
                updateResultDisplay(resultText, 'success');
            });
        }
    } catch (error) {
        console.error('模型训练或加载出错:', error);
        updateResultDisplay('模型训练出错：' + error.message, 'error');
    }
}

/**
 * 更新结果显示区域的函数
 * @param {string} text - 要显示的文本内容
 * @param {string} type - 消息类型：'success', 'error', 'info'
 */
function updateResultDisplay(text, type = 'info') {
    const resultDiv = document.getElementById('result');
    if (resultDiv) {
        // 设置不同类型的样式
        let bgColor = '#f5f5f5';
        if (type === 'success') bgColor = '#d4edda';
        if (type === 'error') bgColor = '#f8d7da';
        
        resultDiv.style.backgroundColor = bgColor;
        resultDiv.style.padding = '15px';
        resultDiv.style.borderRadius = '4px';
        resultDiv.style.border = '1px solid #ddd';
        
        // 更新内容
        resultDiv.innerHTML = text;
    }
}
